package errorgroup

import (
	"context"
	"fmt"
	"sync"
)

type Handler func(ctx context.Context) error
type token struct{}

type Group struct {
	ctx    context.Context
	wg     sync.WaitGroup
	err    error
	once   sync.Once
	tokens chan token
	cancel context.CancelFunc
}

func New(ctx context.Context, limit int) *Group {
	if ctx == nil {
		ctx = context.Background()
	}
	var cancel context.CancelFunc
	ctx, cancel = context.WithCancel(ctx)
	g := &Group{
		ctx:    ctx,
		wg:     sync.WaitGroup{},
		err:    nil,
		once:   sync.Once{},
		tokens: nil,
		cancel: cancel,
	}
	_ = g.SetLimit(limit)
	return g
}

func (g *Group) SetLimit(limit int) error {
	if limit <= 0 {
		g.tokens = nil
	}

	if len(g.tokens) != 0 {
		return fmt.Errorf("errgroup: modify limit while %v goroutines in the group are still active", len(g.tokens))
	}

	g.tokens = make(chan token, limit)
	return nil
}

func (g *Group) done() {
	g.wg.Done()
	if g.tokens != nil {
		<-g.tokens
	}
}

func (g *Group) Do(handler Handler) {
	select {
	case <-g.ctx.Done():
		return
	default:
		if g.tokens != nil {
			g.tokens <- token{}
		}
	}

	g.wg.Add(1)
	go func() {
		defer g.done()

		if err := handler(g.ctx); err != nil {
			g.once.Do(func() {
				g.err = err
				g.cancel()
			})
		}
	}()
}

func (g *Group) Wait() error {
	g.wg.Wait()
	g.cancel()
	return g.err
}
